package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
	
	"github.com/shirou/gopsutil/v3/cpu"
	"github.com/shirou/gopsutil/v3/disk"
	"github.com/shirou/gopsutil/v3/host"
	"github.com/shirou/gopsutil/v3/mem"
)

type Worker struct {
	ID            string
	Name          string
	IPAddress     string
	ControllerURL string
	Registered    bool
	FIOVersion    string
	FIOServerPID  int
	FIOStatus     string // fio服务状态: "running", "stopped", "error"
	FIOServerPort string // fio server监听端口
}

type TaskInfo struct {
	ID        string      `json:"id"`
	TestCaseID string     `json:"test_case_id"`
	FIOConfig interface{} `json:"fio_config"`
}

type TaskResult struct {
	Status  string      `json:"status"`
	JobName string      `json:"jobname,omitempty"`
	Read    interface{} `json:"read,omitempty"`
	Write   interface{} `json:"write,omitempty"`
	Latency interface{} `json:"latency,omitempty"`
	Error   string      `json:"error,omitempty"`
}

func NewWorker(name, controllerURL string) *Worker {
	// 自动获取本地IP地址
	localIP := getLocalIP()
	
	worker := &Worker{
		Name:          name,
		IPAddress:     localIP,
		ControllerURL: controllerURL,
		Registered:    false,
		FIOVersion:    getFIOVersion(),
		FIOStatus:     "unknown", // 初始状态未知
	}
	return worker
}

// getLocalIP 获取本地IP地址
func getLocalIP() string {
	// 首先尝试使用 gopsutil 获取网络接口信息
	interfaces, err := net.Interfaces()
	if err != nil {
		log.Printf("Error getting network interfaces: %v", err)
		return "127.0.0.1"
	}
	
	// 遍历网络接口寻找有效的IPv4地址
	for _, iface := range interfaces {
		// 跳过回环接口和禁用的接口
		if iface.Flags&net.FlagLoopback != 0 || iface.Flags&net.FlagUp == 0 {
			continue
		}
		
		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}
		
		// 遍历接口的地址
		for _, addr := range addrs {
			if ipnet, ok := addr.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
				// 检查是否是IPv4地址
				if ipnet.IP.To4() != nil {
					// 排除私有地址范围外的地址（可根据需要调整）
					ip := ipnet.IP.String()
					if isValidIPAddress(ip) {
						log.Printf("Using local IP address: %s from interface: %s", ip, iface.Name)
						return ip
					}
				}
			}
		}
	}
	
	// 如果找不到有效的外部IP，尝试使用其他方法
	log.Printf("Could not find external IP, trying fallback methods...")
	
	// 回退到 127.0.0.1
	return "127.0.0.1"
}

// isValidIPAddress 检查IP地址是否有效（可以根据需要调整规则）
func isValidIPAddress(ip string) bool {
	// 排除一些常见的无效地址范围
	if ip == "127.0.0.1" || ip == "0.0.0.0" {
		return false
	}
	
	// 可以添加更多的验证规则
	parsedIP := net.ParseIP(ip)
	if parsedIP == nil {
		return false
	}
	
	// 检查是否是私有地址范围（如果需要排除私有地址，可以修改此逻辑）
	// _, private24BitBlock, _ := net.ParseCIDR("10.0.0.0/8")
	// _, private20BitBlock, _ := net.ParseCIDR("172.16.0.0/12")
	// _, private16BitBlock, _ := net.ParseCIDR("192.168.0.0/16")
	
	// 如果你想排除私有地址，可以取消下面的注释
	// if private24BitBlock.Contains(parsedIP) || private20BitBlock.Contains(parsedIP) || private16BitBlock.Contains(parsedIP) {
	// 	return false
	// }
	
	return true
}

// getFIOVersion 获取 fio 版本
func getFIOVersion() string {
	cmd := exec.Command("fio", "--version")
	output, err := cmd.Output()
	if err != nil {
		log.Printf("Error getting fio version: %v", err)
		return "unknown"
	}
	
	// 解析 fio --version 输出，通常格式是 "fio-3.35" 或类似
	versionStr := string(output)
	if strings.HasPrefix(versionStr, "fio-") {
		return strings.TrimSpace(versionStr)
	}
	
	// 如果格式不同，则尝试提取版本号
	lines := strings.Split(versionStr, "\\n")
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "fio-") {
			return trimmed
		}
	}
	
	return "unknown"
}

func (w *Worker) RegisterWithController() error {
	registrationData := map[string]interface{}{
		"name":        w.Name,
		"ip_address":  w.IPAddress,
		"fio_version": w.FIOVersion,  // 添加 fio 版本信息
	}

	jsonData, err := json.Marshal(registrationData)
	if err != nil {
		return err
	}

	resp, err := http.Post(
		fmt.Sprintf("%s/api/workers/register", w.ControllerURL),
		"application/json",
		bytes.NewBuffer(jsonData),
	)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("registration failed with status %d: %s", resp.StatusCode, string(body))
	}

	var workerInfo map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&workerInfo); err != nil {
		return err
	}

	w.ID = workerInfo["id"].(string)
	w.Registered = true

	log.Printf("Worker registered successfully with ID: %s", w.ID)
	return nil
}

func (w *Worker) SendHeartbeat() error {
	if !w.Registered || w.ID == "" {
		return fmt.Errorf("worker not registered")
	}

	// 获取真实的系统和磁盘信息
	var diskInfo []map[string]interface{}

	// 获取CPU信息
	cpuInfo, err := cpu.Info()
	var cpuModel string
	var cpuCores int32
	if err == nil && len(cpuInfo) > 0 {
		cpuModel = cpuInfo[0].ModelName
		cpuCores = cpuInfo[0].Cores
	} else {
		cpuModel = "Unknown"
		cpuCores = 0
	}

	// 获取内存信息
	vmStat, err := mem.VirtualMemory()
	var memoryTotal, memoryUsed, memoryFree uint64
	var memoryPercent float64
	if err == nil {
		memoryTotal = vmStat.Total
		memoryUsed = vmStat.Used
		memoryFree = vmStat.Free
		memoryPercent = vmStat.UsedPercent
	} else {
		memoryTotal = 0
		memoryUsed = 0
		memoryFree = 0
		memoryPercent = 0
	}

	// 获取主机信息
	hostInfo, err := host.Info()
	var osInfo, kernelInfo, arch string
	var uptime uint64
	if err == nil {
		osInfo = hostInfo.OS
		kernelInfo = hostInfo.KernelVersion
		arch = hostInfo.PlatformFamily
		if arch == "" {
			arch = hostInfo.KernelArch
		}
		uptime = hostInfo.Uptime
	} else {
		osInfo = "Unknown"
		kernelInfo = "Unknown"
		arch = "Unknown"
		uptime = 0
	}

	// 获取磁盘信息
	partitions, err := disk.Partitions(false)
	if err != nil {
		log.Printf("Error getting disk partitions: %v", err)
	}

	for _, partition := range partitions {
		usage, err := disk.Usage(partition.Mountpoint)
		if err == nil {
			diskInfo = append(diskInfo, map[string]interface{}{
				"name":         partition.Device,
				"mount_point":  partition.Mountpoint,
				"total_space":  usage.Total,
				"used_space":   usage.Used,
				"free_space":   usage.Free,
				"used_percent": usage.UsedPercent,
				"file_system":  partition.Fstype,
			})
		} else {
			log.Printf("Error getting disk usage for %s: %v", partition.Mountpoint, err)
		}
	}

	// 构建心跳数据，包含真实的系统和磁盘信息
	heartbeatData := map[string]interface{}{
		"system_info": map[string]interface{}{
			"cpu_info":       cpuModel,
			"cpu_cores":      cpuCores,
			"memory_total":   memoryTotal,
			"memory_used":    memoryUsed,
			"memory_free":    memoryFree,
			"memory_percent": memoryPercent,
			"os":             osInfo,
			"kernel":         kernelInfo,
			"arch":           arch,
			"uptime":         uptime,
		},
		"disk_info": diskInfo,
		"fio_version": w.FIOVersion,      // 添加 fio 版本信息
		"fio_status":  w.FIOStatus,       // 添加 fio 状态信息
		"fio_port":    w.FIOServerPort,   // 添加 fio 端口信息
	}

	// 将心跳数据转换为JSON
	jsonData, err := json.Marshal(heartbeatData)
	if err != nil {
		return err
	}

	resp, err := http.Post(
		fmt.Sprintf("%s/api/workers/%s/heartbeat", w.ControllerURL, w.ID),
		"application/json",
		bytes.NewBuffer(jsonData),
	)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("heartbeat failed with status %d: %s", resp.StatusCode, string(body))
	}

	log.Printf("Heartbeat sent successfully")
	return nil
}

func (w *Worker) GetPendingTasks() ([]TaskInfo, error) {
	if !w.Registered || w.ID == "" {
		return nil, fmt.Errorf("worker not registered")
	}

	resp, err := http.Get(fmt.Sprintf("%s/api/tasks/worker/%s/pending", w.ControllerURL, w.ID))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("failed to get pending tasks: %d - %s", resp.StatusCode, string(body))
	}

	var tasks []TaskInfo
	if err := json.NewDecoder(resp.Body).Decode(&tasks); err != nil {
		return nil, err
	}

	return tasks, nil
}

func (w *Worker) ExecuteFIOTask(taskID string, fioConfig interface{}) error {
	log.Printf("Executing FIO task %s", taskID)

	// 将FIO配置转换为INI格式
	configINI, err := w.convertFIOConfigToINI(fioConfig)
	if err != nil {
		log.Printf("Error converting FIO config: %v", err)
		result := TaskResult{
			Status: "failed",
			Error:  err.Error(),
		}
		w.ReportTaskCompletion(taskID, result)
		return err
	}

	// 创建临时配置文件
	configFile := fmt.Sprintf("/tmp/fio_config_%s.fio", taskID)
	if err := w.writeFIOConfigFile(configFile, configINI); err != nil {
		log.Printf("Error writing config file: %v", err)
		result := TaskResult{
			Status: "failed",
			Error:  err.Error(),
		}
		w.ReportTaskCompletion(taskID, result)
		return err
	}

	// 使用 fio client 连接到本地 fio server 执行测试
	// 这是真正的分布式 fio 测试模式
	cmd := exec.Command("fio", fmt.Sprintf("--client=%s", w.IPAddress), configFile)

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err = cmd.Run()

	// 清理临时配置文件
	defer func() {
		exec.Command("rm", configFile).Run()
	}()

	if err != nil {
		log.Printf("FIO client execution failed: %v", err)
		// 如果client模式失败，回退到本地执行模式
		log.Printf("Falling back to local FIO execution")
		cmdLocal := exec.Command("fio", configFile)
		var stdoutLocal, stderrLocal bytes.Buffer
		cmdLocal.Stdout = &stdoutLocal
		cmdLocal.Stderr = &stderrLocal
		errLocal := cmdLocal.Run()
		if errLocal != nil {
			result := TaskResult{
				Status: "failed",
				Error:  fmt.Sprintf("FIO execution failed in both client and local mode: %v, Local err: %v\nStderr: %s", err, errLocal, stderrLocal.String()),
			}

			w.ReportTaskCompletion(taskID, result)
			return errLocal
		}
		// 解析本地执行结果
		result := w.parseFIOOutput(stdoutLocal.String())
		result.Status = "completed"
		log.Printf("FIO task %s completed using local execution", taskID)
		w.ReportTaskCompletion(taskID, result)
		return nil
	}

	// 解析FIO输出结果
	result := w.parseFIOOutput(stdout.String())
	result.Status = "completed"

	log.Printf("FIO task %s completed using client/server mode", taskID)
	w.ReportTaskCompletion(taskID, result)

	return nil
}

func (w *Worker) convertFIOConfigToINI(fioConfig interface{}) (string, error) {
	// 这是一个简化版本，实际实现可能需要更复杂的配置转换
	configMap, ok := fioConfig.(map[string]interface{})
	if !ok {
		return "", fmt.Errorf("invalid FIO config format")
	}

	var iniBuilder strings.Builder
	iniBuilder.WriteString("[global]\n")

	// 转换全局配置参数
	for key, value := range configMap {
		if key != "jobs" { // jobs需要特别处理
			switch v := value.(type) {
			case string:
				iniBuilder.WriteString(fmt.Sprintf("%s=%s\n", key, v))
			case int:
				iniBuilder.WriteString(fmt.Sprintf("%s=%d\n", key, v))
			case float64:
				iniBuilder.WriteString(fmt.Sprintf("%s=%.0f\n", key, v))
			case bool:
				iniBuilder.WriteString(fmt.Sprintf("%s=%t\n", key, v))
			default:
				iniBuilder.WriteString(fmt.Sprintf("%s=%v\n", key, v))
			}
		}
	}

	// 如果有具体的jobs配置
	if jobs, exists := configMap["jobs"]; exists {
		jobsArray, ok := jobs.([]interface{})
		if ok {
			for i, job := range jobsArray {
				jobMap, ok := job.(map[string]interface{})
				if !ok {
					continue
				}

				jobName, hasName := jobMap["name"]
				if !hasName {
					jobName = fmt.Sprintf("job_%d", i)
				}

				iniBuilder.WriteString(fmt.Sprintf("\n[%s]\n", jobName))
				for key, value := range jobMap {
					if key != "name" {
						switch v := value.(type) {
						case string:
							iniBuilder.WriteString(fmt.Sprintf("%s=%s\n", key, v))
						case int:
							iniBuilder.WriteString(fmt.Sprintf("%s=%d\n", key, v))
						case float64:
							iniBuilder.WriteString(fmt.Sprintf("%s=%.0f\n", key, v))
						case bool:
							iniBuilder.WriteString(fmt.Sprintf("%s=%t\n", key, v))
						default:
							iniBuilder.WriteString(fmt.Sprintf("%s=%v\n", key, v))
						}
					}
				}
			}
		}
	}

	return iniBuilder.String(), nil
}

func (w *Worker) writeFIOConfigFile(filename, content string) error {
	// 实际写入配置文件到磁盘
	return os.WriteFile(filename, []byte(content), 0644)
}

func (w *Worker) parseFIOOutput(output string) TaskResult {
	// 这是一个简化的FIO输出解析器
	result := TaskResult{
		Read:  make(map[string]interface{}),
		Write: make(map[string]interface{}),
	}

	// 在实际实现中，这里应该解析FIO的JSON输出或其他格式的输出
	// 为了演示，我们仅返回一个基本结果
	result.Read = map[string]interface{}{
		"iops":      10000,
		"bandwidth": 100000,
	}
	result.Write = map[string]interface{}{
		"iops":      5000,
		"bandwidth": 50000,
	}
	result.Latency = map[string]interface{}{
		"avg": 100,
		"max": 500,
	}

	return result
}

func (w *Worker) ReportTaskCompletion(taskID string, result TaskResult) error {
	jsonData, err := json.Marshal(result)
	if err != nil {
		return err
	}

	resp, err := http.Post(
		fmt.Sprintf("%s/api/tasks/%s/complete", w.ControllerURL, taskID),
		"application/json",
		bytes.NewBuffer(jsonData),
	)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("failed to report task completion: %d - %s", resp.StatusCode, string(body))
	}

	log.Printf("Task %s completion reported successfully", taskID)
	return nil
}

func (w *Worker) Run(heartbeatInterval time.Duration) {
	// 启动 fio server
	if err := w.StartFIOServer(); err != nil {
		log.Printf("Warning: Failed to start fio server: %v", err)
	} else {
		log.Printf("fio server started successfully")
	}

	if err := w.RegisterWithController(); err != nil {
		log.Fatalf("Failed to register with controller: %v", err)
	}

	log.Printf("Worker %s running with ID %s", w.Name, w.ID)

	ticker := time.NewTicker(heartbeatInterval)
	defer ticker.Stop()

	// 启动 fio 状态检查 ticker（每10秒检查一次）
	fioCheckTicker := time.NewTicker(10 * time.Second)
	defer fioCheckTicker.Stop()

	for {
		select {
		case <-ticker.C:
			// 发送心跳
			if err := w.SendHeartbeat(); err != nil {
				log.Printf("Error sending heartbeat: %v", err)
			}

			// 获取并执行待处理任务
			tasks, err := w.GetPendingTasks()
			if err != nil {
				log.Printf("Error getting pending tasks: %v", err)
				continue
			}

			for _, task := range tasks {
				log.Printf("Executing task %s", task.ID)
				if err := w.ExecuteFIOTask(task.ID, task.FIOConfig); err != nil {
					log.Printf("Error executing task %s: %v", task.ID, err)
				}
			}
		case <-fioCheckTicker.C:
			// 定期检查 fio 状态
			w.checkFIOStatus()
		}
	}
}

// StartFIOServer 启动 fio server 模式
func (w *Worker) StartFIOServer() error {
		// 使用 fio 的 server 模式启动服务，默认监听在 8765 端口
	// fio server模式会监听端口等待来自client的测试配置
	// 使用worker的IP地址让其他客户端可以连接
	cmd := exec.Command("fio", fmt.Sprintf("--server=ip:%s,8765", w.IPAddress))  // 启动 fio server 模式，绑定到worker的IP地址和端口
	
	// 重定向标准输出和错误输出到日志文件
	logFile, err := os.Create(fmt.Sprintf("/tmp/fio_server_%d.log", time.Now().Unix()))
	if err != nil {
		log.Printf("Warning: Could not create log file for fio server: %v", err)
	} else {
		cmd.Stdout = logFile
		cmd.Stderr = logFile
	}
	
	if err := cmd.Start(); err != nil {
		w.FIOStatus = "error"
		return fmt.Errorf("failed to start fio server: %v", err)
	}
	
	w.FIOServerPID = cmd.Process.Pid
	w.FIOStatus = "running"
	w.FIOServerPort = "8765"  // 设置默认端口
	log.Printf("fio server started with PID: %d on port: %s", w.FIOServerPID, w.FIOServerPort)
	
	// 启动一个 goroutine 来监控 fio server 进程
	go func() {
		err := cmd.Wait()
		if err != nil {
			log.Printf("fio server exited with error: %v", err)
			w.FIOStatus = "error"
		} else {
			log.Printf("fio server exited normally")
			w.FIOStatus = "stopped"
		}
		w.FIOServerPID = 0 // 清除 PID
		w.FIOServerPort = "" // 清除端口
		if logFile != nil {
			logFile.Close()
		}
	}()
	
	return nil
}

// checkFIOStatus 检查 fio 服务状态
func (w *Worker) checkFIOStatus() {
	// 如果没有启动 fio server，直接返回
	if w.FIOServerPID <= 0 {
		w.FIOStatus = "stopped"
		w.FIOServerPort = ""
		return
	}

	// 检查进程是否还在运行
	process, err := os.FindProcess(w.FIOServerPID)
	if err != nil {
		w.FIOStatus = "error"
		log.Printf("Error finding fio process %d: %v", w.FIOServerPID, err)
		return
	}

	// 在 Unix/Linux 系统上，信号0不会真正发送信号，但会检查进程是否存在
	err = process.Signal(syscall.Signal(0))
	if err != nil {
		// 进程不存在或无法访问
		w.FIOStatus = "stopped"
		w.FIOServerPID = 0
		w.FIOServerPort = ""
		log.Printf("fio process %d is not running", w.FIOServerPID)
		return
	}

	// 进程存在，尝试连接到 fio server 来确认服务状态
	// 如果已经有端口信息，先尝试该端口
	var connectedPort string
	if w.FIOServerPort != "" {
		// 尝试已知端口，使用worker的IP地址
		conn, err := net.DialTimeout("tcp", w.IPAddress+":"+w.FIOServerPort, 2*time.Second)
		if err == nil {
			conn.Close()
			// 端口可用，确认服务运行
			w.FIOStatus = "running"
			log.Printf("fio server is running and accessible on %s:%s (PID: %d)", w.IPAddress, w.FIOServerPort, w.FIOServerPID)
			return
		}
	}

	// 如果已知端口不可用或未知，尝试常见端口
	ports := []string{"8765", "8888", "8766"}
	connected := false
	for _, port := range ports {
		conn, err := net.DialTimeout("tcp", w.IPAddress+":"+port, 1*time.Second)
		if err == nil {
			connected = true
			connectedPort = port
			conn.Close()
			break
		}
	}
	
	if !connected {
		w.FIOStatus = "error"
		w.FIOServerPort = ""
		log.Printf("Cannot connect to fio server on any common port: %v", err)
		return
	}

	// 成功连接到 fio server，更新端口信息
	w.FIOStatus = "running"
	w.FIOServerPort = connectedPort
	log.Printf("fio server is running and accessible on port %s (PID: %d)", w.FIOServerPort, w.FIOServerPID)
}

func main() {
	worker := NewWorker("worker-01", "http://localhost:8091")  // Use the correct controller port
	worker.Run(30 * time.Second)
}